import json
import subprocess
from math import log

from file_generator import generate_alphabet, random_string

ALPHABET_SIZES = [6, 7, 8]
ALGORITHMS = ['Brute force', 'Horspool', 'BNDM', 'BOM']
COMMAND = './search {alphabet_size} ~/y/ri-data/random{alphabet_size}.txt {patterns}'
MAX_LENGTH = 256
NUM_SAMPLES = 5
VERBOSE = True
SAVE_FILE = 'results-Alphabet{0}.json'


def load(n):
    with open(SAVE_FILE.format(n)) as fp:
        return json.load(fp)


def save(data, n):
    with open(SAVE_FILE.format(n), 'w') as fp:
        json.dump(data, fp)


def custom_range(i, n):
    while i <= n:
        yield i
        i = i + 2 ** int(log(i, 4) if i > 0 else 0)


def interpolate(vector):
    return vector


def generate_test_cases(alphabet, max_length, num_samples):
    """
    Generates a list of sets.
    Each set contains `num_samples` of strings of equal length.
    Lengths in each set range from 1 to `max_length`.
    """
    test_set = []
    for n in custom_range(1, max_length):
        strings = []
        for i in range(num_samples):
            strings.append(random_string(n, alphabet))
        test_set.append(strings)
    return test_set


def test_cases_to_list(test_cases):
    """
    Returns a list with all the strings that compose the test set.
    """
    return [string for string_set in test_cases for string in string_set]


def run_experiments(algorithms, alphabet_size, test_cases):
    """
    Takes a test case and executes it for each algorithm.
    """
    # Run C program passing all the strings in the test case
    command = COMMAND.format(
        alphabet_size=alphabet_size,
        patterns=' '.join(test_cases_to_list(test_cases)))
    if VERBOSE:
        print(command)
    raw = subprocess.check_output(command, shell=True)
    data = dict()
    output = str(raw, 'utf-8')
    if VERBOSE:
        print(output)
    p = 0
    for algo in algorithms:
        # Process output, a list with the result for each string
        lengths = []
        counts = []
        times = []
        results = output.split('\n')
        for string_set in test_cases:
            total_count = 0
            total_time = 0
            for string in string_set:
                pattern, count, time = results[p].split(',')
                p += 1
                total_count += int(count)
                total_time += float(time)
            average_count = total_count / len(string_set)
            average_time = total_time / len(string_set)
            lengths.append(len(pattern))
            counts.append(average_count)
            times.append(average_time)
        data[algo] = (lengths, counts, times)
        if (VERBOSE):
            print("total time: {0}".format(sum(times) * NUM_SAMPLES))
    return data


def generate_figures(title, data):
    """
    Generates some .png graphs using data from the experiment.
    """
    try:
        from matplotlib import pyplot as plt
    except:
        print("matplotlib.pyplot required to generate the figures")
    else:
        # Time & Pattern length (for each algorithm)
        names = []
        for alg, results in data.items():
            names.append(alg)
            lengths, counts, times = results
            plt.plot(lengths, interpolate(times))
        plt.title(title)
        plt.xlabel('Pattern length')
        plt.ylabel('Execution time')
        plt.legend(names)
        plt.savefig('results-{0}-length-time.png'.format(title))
        plt.xscale('log', basex=2)
        plt.savefig('results-{0}-length-time-log.png'.format(title))
        plt.clf()
        # Count & Pattern length (the same for each algorithm)
        lengths, counts, times = data[ALGORITHMS[0]]
        plt.plot(lengths, counts)
        plt.title(title)
        plt.xlabel('Pattern length')
        plt.ylabel('Times found')
        plt.savefig('results-{0}-length-count.png'.format(title))
        plt.xscale('log', basex=2)
        plt.savefig('results-{0}-length-count-log.png'.format(title))
        plt.clf()
        bests = []
        for i, length in enumerate(lengths):
            time = 9999
            best = -1
            for k, alg in enumerate(ALGORITHMS):
                if data[alg][2][i] < time:
                    time = data[alg][2][i]
                    best = k
            bests.append(best)
        plt.plot(lengths, bests, marker='o', ls='', color='black')
        plt.yticks([0, 1, 2, 3], ALGORITHMS)
        plt.title(title)
        plt.xlabel('Pattern length')
        plt.ylabel('Best algorithm')
        plt.savefig('results-{0}-best.png'.format(title))
        plt.xscale('log', basex=2)
        plt.savefig('results-{0}-best-log.png'.format(title))
        plt.clf()

if __name__ == '__main__':
    for n in ALPHABET_SIZES:
        title = 'Alphabet{0}'.format(n)
        # results = load(n)
        alphabet = generate_alphabet(n)
        test_cases = generate_test_cases(alphabet, MAX_LENGTH, NUM_SAMPLES)
        results = run_experiments(ALGORITHMS, n, test_cases)
        save(results, n)
        generate_figures(title, results)
